#include "data_structures.hpp"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#include <vector>
#include <memory>
#include <algorithm>
#include <cmath>
#include <limits>

namespace EddyDetection {

    namespace {
        // Haversine Spherical distance calculation (unit: kilometers)
        double haversine_distance_km(const Eigen::Vector2d& p1, const Eigen::Vector2d& p2) {
            const double R = 6371.0;
            double lat1 = p1[0] * M_PI / 180.0;
            double lon1 = p1[1] * M_PI / 180.0;
            double lat2 = p2[0] * M_PI / 180.0;
            double lon2 = p2[1] * M_PI / 180.0;

            double dlon = lon2 - lon1;
            double dlat = lat2 - lat1;

            double a = std::sin(dlat / 2.0) * std::sin(dlat / 2.0) +
                std::cos(lat1) * std::cos(lat2) *
                std::sin(dlon / 2.0) * std::sin(dlon / 2.0);
            return R * 2.0 * std::atan2(std::sqrt(a), std::sqrt(1.0 - a));
        }

        // Octree node
        struct OctreeNode {
            Eigen::Vector3d min_bound; // Minimum boundary (lat, lon, depth)
            Eigen::Vector3d max_bound; // Maximum boundary
            std::vector<const EddyContour*> eddies; // Eddy within the current node
            std::array<std::unique_ptr<OctreeNode>, 8> children; // 8 child nodes

            OctreeNode(const Eigen::Vector3d& min, const Eigen::Vector3d& max)
                : min_bound(min), max_bound(max) {
            }

            // Insert Eddy
            void insert(const EddyContour* eddy, int max_depth = 5, int depth = 0) {
                Eigen::Vector3d point(eddy->center.x(), eddy->center.y(), eddy->depth);
                if (!contains(point)) return;

                if (depth >= max_depth || eddies.size() < 10) {
                    eddies.push_back(eddy);
                    return;
                }

                if (children[0] == nullptr) split();
                for (auto& child : children) {
                    child->insert(eddy, max_depth, depth + 1);
                }
            }

            // Split Node
            void split() {
                Eigen::Vector3d mid = (min_bound + max_bound) * 0.5;
                children[0] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(min_bound.x(), min_bound.y(), min_bound.z()),
                    Eigen::Vector3d(mid.x(), mid.y(), mid.z()));
                children[1] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(mid.x(), min_bound.y(), min_bound.z()),
                    Eigen::Vector3d(max_bound.x(), mid.y(), mid.z()));
                children[2] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(min_bound.x(), mid.y(), min_bound.z()),
                    Eigen::Vector3d(mid.x(), max_bound.y(), mid.z()));
                children[3] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(mid.x(), mid.y(), min_bound.z()),
                    Eigen::Vector3d(max_bound.x(), max_bound.y(), mid.z()));
                children[4] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(min_bound.x(), min_bound.y(), mid.z()),
                    Eigen::Vector3d(mid.x(), mid.y(), max_bound.z()));
                children[5] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(mid.x(), min_bound.y(), mid.z()),
                    Eigen::Vector3d(max_bound.x(), mid.y(), max_bound.z()));
                children[6] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(min_bound.x(), mid.y(), mid.z()),
                    Eigen::Vector3d(mid.x(), max_bound.y(), max_bound.z()));
                children[7] = std::make_unique<OctreeNode>(
                    Eigen::Vector3d(mid.x(), mid.y(), mid.z()),
                    Eigen::Vector3d(max_bound.x(), max_bound.y(), max_bound.z()));
            }

            bool contains(const Eigen::Vector3d& point) const {
                return point.x() >= min_bound.x() && point.x() <= max_bound.x() &&
                    point.y() >= min_bound.y() && point.y() <= max_bound.y() &&
                    point.z() >= min_bound.z() && point.z() <= max_bound.z();
            }

            void query_spherical(//************************3
                const Eigen::Vector2d& center_xy, double depth_z, double radius_km,
                const EddyContourDataCyclone eddy_type, const std::vector<const EddyContour*>& exclude,
                std::vector<const EddyContour*>& result) const {

                Eigen::Vector3d node_center = (min_bound + max_bound) * 0.5;
                double approx_dist_xy = std::hypot(
                    node_center.x() - center_xy.x(),
                    node_center.y() - center_xy.y());
                //double approx_dist_z = std::abs(node_center.z() - depth_z); // this is not used

                if (approx_dist_xy * 111.32 > radius_km * 1.5) { 
                    std::cout << "approx_dist_xy * 111.32 > radius_km * 1.5, " << approx_dist_xy * 111.32 <<">" << radius_km * 1.5 << std::endl;
                    return;
                }else {
                    //std::cout << node_center.x() << ", " << node_center.y() << ", " << center_xy.x() << ", " << center_xy.y() << std::endl;
                    //std::cout << "There is approx_dist_xy * 111.32 <= radius_km * 1.5!!!!!!!!!!!!!!!!!!!!!!!!" << std::endl;
                    std::cout << "eddies.size = " << eddies.size() << std::endl;
                }

                for (const auto& eddy : eddies) {
                    if (eddy->type != eddy_type) {
                        //std::cout << "type is not ok " << std::endl;
                        continue;
                    } 
                    if (std::find(exclude.begin(), exclude.end(), eddy) != exclude.end()) {
                        //std::cout << "exclude is not ok " << std::endl;
                        continue;
                    }
                        
                    double dist = haversine_distance_km(center_xy, eddy->center);
                    //std::cout << "center_xy.x, center_xy.y, center.x, center.y = " << center_xy.x() << ", " << center_xy.y() << "; " << eddy->center.x() << ", " << eddy->center.y() << std::endl;
                    if (dist <= radius_km) {
                        std::cout << "result.push_back(eddy)" << std::endl;
                        result.push_back(eddy);
                    }else
                    {
                        std::cout << "dist > radius_km:  " << dist <<" > "<<radius_km << std::endl;
                    }
                }

                for (const auto& child : children) {
                    if (child) {
                        child->query_spherical(center_xy, depth_z, radius_km, eddy_type, exclude, result);
                    }
                }
            }
        };

        class EddyOctree {
            std::unique_ptr<OctreeNode> root;
            double max_distance_km;

        public:
            EddyOctree(const std::vector<const EddyContour*>& eddies, double max_dist_km)
                : max_distance_km(max_dist_km) {

                Eigen::Vector3d min_bound(90, 180, std::numeric_limits<double>::max());
                Eigen::Vector3d max_bound(-90, -180, -std::numeric_limits<double>::max());

                // Assign a value to the bound lat&lon of an Octree
                for (const auto& eddy : eddies) {
                    min_bound.x() = std::min(min_bound.x(), eddy->center.x());
                    min_bound.y() = std::min(min_bound.y(), eddy->center.y());
                    min_bound.z() = std::min(min_bound.z(), eddy->depth);
                    max_bound.x() = std::max(max_bound.x(), eddy->center.x());
                    max_bound.y() = std::max(max_bound.y(), eddy->center.y());
                    max_bound.z() = std::max(max_bound.z(), eddy->depth);
                }
                root = std::make_unique<OctreeNode>(min_bound, max_bound);

                for (const auto& eddy : eddies) {
                    root->insert(eddy);
                }
            }

            std::vector<const EddyContour*> find_nearby(//*******************************1
                const EddyContour* eddy, const std::vector<const EddyContour*>& exclude) {
                std::vector<const EddyContour*> result;
                root->query_spherical(//***************************2
                    eddy->center, eddy->depth, max_distance_km,
                    eddy->type, exclude, result);
                return result;
            }
        };
    }

    // Main Logic Function
    std::vector<Eddy3D> ThreeDLinker::link3DEddies(
        const std::vector<std::vector<EddyContour>>& all_layers_eddies,
        const EddyDetectionConfig& config) {

        // Flatten all Eddy pointers
        std::vector<const EddyContour*> all_eddies;
        for (size_t layer_idx = 0; layer_idx < all_layers_eddies.size(); ++layer_idx) {
            for (const auto& eddy : all_layers_eddies[layer_idx]) {
                all_eddies.push_back(&eddy);
            }
        }

        // Building an Octree
        EddyOctree octree(all_eddies, config.max_layer_distance_km);

        std::vector<Eddy3D> final_eddies;
        std::vector<const EddyContour*> used_eddies;

        for (const auto& eddy_start : all_eddies) {
            if (std::find(used_eddies.begin(), used_eddies.end(), eddy_start) != used_eddies.end()) {
                continue;
            }

            std::vector<const EddyContour*> chain = { eddy_start };
            used_eddies.push_back(eddy_start);

            // Track with a variable pointer
            const EddyContour* eddy_ptr = eddy_start;

            // Search up for adjacent layers
            for (size_t next_layer = eddy_ptr->layer_index + 1;
                next_layer < all_layers_eddies.size(); ++next_layer) {

                auto candidates = octree.find_nearby(eddy_ptr, used_eddies);//***************0

                std::cout << "Layer " << eddy_ptr->layer_index << " -> " << next_layer 
                    << ": found" << candidates.size() << " candidates" << std::endl;

                if (candidates.empty()) {
                    std::cout << "No candidates found for eddy at (" << eddy_ptr->center.x() << ", " <<eddy_ptr->center.y()
                        << ") depth="<<eddy_ptr->depth << std::endl;
                    break;
                }
                // Select the nearest candidate Eddy
                const EddyContour* best_match = nullptr;
                double min_dist = std::numeric_limits<double>::max();
                for (const auto& candidate : candidates) {
                    double dist = haversine_distance_km(eddy_ptr->center, candidate->center);
                    if (dist < min_dist) {
                        min_dist = dist;
                        best_match = candidate;
                    }
                }

                std::cout << "There is a best_dist = " << min_dist << std::endl;

                if (best_match && min_dist <= config.max_layer_distance_km) {
                    chain.push_back(best_match);
                    used_eddies.push_back(best_match);
                    eddy_ptr = best_match;

                    //std::cout << "There is a min_dist <= config.max_layer_distance_km " << std::endl;
                }
                else {
                    std::cout << "There is a min_dist = " << min_dist << std::endl;
                    break;
                }
            }
            //std::cout << "chain.size() = " << chain.size() << std::endl;
            std::cout << std::endl;

            if (chain.size() >= config.min_vertical_layers) {
                final_eddies.push_back(create3DEddy(chain));
                std::cout << "chain.size() = " << chain.size() << std::endl;
            }
        }

        return final_eddies;
    }

    Eddy3D ThreeDLinker::create3DEddy(const std::vector<const EddyContour*>& chain) {
        Eddy3D eddy_3d;
        eddy_3d.type = chain[0]->type;
        eddy_3d.layer_count = chain.size();
        eddy_3d.layers.reserve(chain.size());

        double sum_lat = 0, sum_lon = 0, sum_depth = 0;
        double min_depth = chain[0]->depth;
        double max_depth = chain[0]->depth;
        double total_volume = 0;

        for (size_t i = 0; i < chain.size(); ++i) {
            eddy_3d.layers.push_back(*chain[i]);
            sum_lat += chain[i]->center.x();
            sum_lon += chain[i]->center.y();
            sum_depth += chain[i]->depth;
            min_depth = std::min(min_depth, chain[i]->depth);
            max_depth = std::max(max_depth, chain[i]->depth);

            if (i > 0) {
                double depth_diff = std::abs(chain[i]->depth - chain[i - 1]->depth) / 1000.0; // Convert to kilometers
                double avg_area = (chain[i]->area_km2 + chain[i - 1]->area_km2) / 2.0;
                total_volume += avg_area * depth_diff;
            }
        }

        eddy_3d.center_3d = {
            sum_lat / chain.size(),
            sum_lon / chain.size(),
            sum_depth / chain.size()
        };
        eddy_3d.vertical_extent = { min_depth, max_depth };
        eddy_3d.volume_km3 = total_volume;
        eddy_3d.id = "3D-" + eddy_3d.type.substr(0, 1) + "-" +
            std::to_string(static_cast<int>(min_depth)) + "m-" +
            std::to_string(chain.size()) + "L";

        return eddy_3d;
    }

} // namespace EddyDetection